import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

import time
import json
import argparse
from tqdm import tqdm
from src.llamafactory.chat.chat_model import ChatModel


def get_current_time(format_string="%Y-%m-%d-%H-%M-%S"):
    return time.strftime(format_string, time.localtime())


output = 'saves/qwen2-7b-lora/atom_sft_ep5_bs8_lr2e-4/checkpoint-600'
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument("--model_name_or_path", default="/home/amax/AI/mm/AiCut/pretrain/Qwen2-7B-Instruct", type=str,
                    help="Model name or path")
parser.add_argument("--adapter_name_or_path", default=output, type=str, help="Checkpoint directory")
parser.add_argument("--template", default="qwen", type=str, help="Template type")
parser.add_argument("--finetuning_type", default="lora", type=str, help="Fine-tuning type")
args = parser.parse_args()

chat_model = ChatModel(vars(args))

prompt_template = """你是一名专业的智能体agent专家，请根据可调用的工具库 apis 和用户多轮对话，给出多轮对话中应使用的工具 api 名称和参数。

## 要求
1.用户的多轮对话中可能要调用多个工具的api或者多次调用某个工具的api，给出工具名。
2.根据多轮对话内容，判断出在第几轮可以满足工具调用的参数需求，然后将api名称和参数加入该轮结果。

工具库 apis:
{apis}

多轮对话:
{user_messages}

输出："""


def align(user_messages, targets):
    '''对齐多轮对话和输出
    '''
    if len(user_messages) > 2:
        res = []
        for t in targets:
            if t == []:
                res.append(t)
            else:
                merge_t = []
                for item in t:
                    if isinstance(item, list) and item:
                        print('单轮内存在子列表')
                        merge_t.extend(item)
                res.append(merge_t)
    elif len(user_messages) == 1:
        res = []
        for t in targets:
            if isinstance(t, list) and t:
                res.extend(t)
        res = [res]
    elif len(user_messages) == 2:
        res = []
        for t in targets:
            if isinstance(t, list) and t:
                res.extend(t)
        if targets[0] == []:
            res = [[], res]
        elif targets[-1] == []:
            res = [res, []]
        else:
            print(f'两轮对话, targets: {targets}')
            res = targets
    return res


def fix_json(res):
    '''修正结果 json 格式
    '''
    text = ''
    pair_stack = []
    for c in res:
        if c in '{[':
            pair_stack.append(c)
            text += c
        elif c == ']':
            if pair_stack == []:
                continue
            elif pair_stack[-1] == '{':
                pair_stack.pop()
                text += '}' + c
            elif pair_stack[-1] in '[':
                pair_stack.pop()
                text += c
        elif c == '}':
            if pair_stack == []:
                continue
            elif pair_stack[-1] == '[':
                pair_stack.pop()
                text += ']' + c
            elif pair_stack[-1] in '{':
                pair_stack.pop()
                text += c
        else:
            text += c

    # 最后检查 pair_stack 和 [] 数量
    if pair_stack == ['['] and text.count('[') > text.count(']'):
        text += ']'

    return text


num = 0
history = []
with open('data/合成数据决赛赛题.jsonl','r', encoding="utf-8") as fr, open(f"data/DL_submission_testb_{get_current_time()}.jsonl", "a", encoding="utf-8") as fw:
    for line in tqdm(fr):
        line = line.strip()
        if line == '':
            continue
        d = json.loads(line)

        apis = d['apis']
        user_messages = d['user_messages']
        
        prompt = prompt_template.replace('{apis}', json.dumps(apis, ensure_ascii=False))
        prompt = prompt.replace('{user_messages}', json.dumps(user_messages, ensure_ascii=False))
        messages = [{"role": "user", "content": prompt}]

        response = chat_model.chat(
            messages,
            history=[],
            do_sample=False,
            max_new_tokens=4096,
            top_p=0.8,
            temperature=0.1,
            num_beams=3,
        )
        res = response[0].response_text.strip()
        res = res.replace('true', 'True').replace('false', 'False')
        res = res.replace('}, [', '}], [')
        
        try:
            targets = eval(res) # list
        except Exception as error:
            print(f'{error}\n转换 list 格式失败：{res}')
            try:
                res = fix_json(res)
                targets = eval(res) # list
                print('修正成功')
            except:
                targets = []
                print('修正失败')
        
        try:
            if len(user_messages) != len(targets) and len(targets) != 0:
                targets = align(user_messages, targets)
                num += 1
                print(f'共对齐{num}个结果')
            d['targets'] = targets
            fw.write(json.dumps(d, ensure_ascii=False) + "\n")
        except:
            d['targets'] = []
            fw.write(json.dumps(d, ensure_ascii=False) + "\n")